/*
 * Copyright 2016-2022 dromara.org.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.dromara.mendmix.gateway.filter.subhandler;

import java.math.BigDecimal;
import java.math.RoundingMode;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.regex.Pattern;

import org.apache.commons.lang3.StringUtils;
import org.dromara.mendmix.cache.CacheExpires;
import org.dromara.mendmix.cache.CacheUtils;
import org.dromara.mendmix.common.GlobalContext;
import org.dromara.mendmix.common.MendmixBaseException;
import org.dromara.mendmix.common.exception.MainErrorType;
import org.dromara.mendmix.common.model.ApiInfo;
import org.dromara.mendmix.common.util.DigestUtils;
import org.dromara.mendmix.common.util.ResourceUtils;
import org.dromara.mendmix.gateway.GatewayConstants;
import org.dromara.mendmix.gateway.ProxySystemInfoHub;
import org.dromara.mendmix.gateway.api.SystemApi;
import org.dromara.mendmix.gateway.component.ratelimit.LocalRatelimiter;
import org.dromara.mendmix.gateway.component.ratelimit.RateLimiter;
import org.dromara.mendmix.gateway.component.ratelimit.RedisRatelimiter;
import org.dromara.mendmix.gateway.filter.SubFilterHandler;
import org.dromara.mendmix.gateway.helper.RequestContextHelper;
import org.dromara.mendmix.gateway.model.BizSystem;
import org.dromara.mendmix.gateway.model.BizSystemModule;
import org.dromara.mendmix.gateway.model.RatelimitStrategy;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.cloud.client.discovery.DiscoveryClient;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.http.server.reactive.ServerHttpRequest.Builder;
import org.springframework.web.server.ServerWebExchange;

/**
 * @description <br>
 * @author <a href="mailto:vakinge@gmail.com">jiangwei</a>
 * @date 2025年5月18日
 */
public class RatelimitHandler implements SubFilterHandler {

	private static Logger logger = LoggerFactory.getLogger("org.dromara.mendmix");
	private static final String ANY_URI_WILDCARD = "/*";
	private static final String POD_NUM_CACHE_KEY = "podNums:" + GlobalContext.APPID;

	// <serviceId,<url,rule>>
	private Map<String, Map<String, RatelimitStrategy>> serviceStrategies = new ConcurrentHashMap<>();
	private Map<String, Map<Pattern, RatelimitStrategy>> wildcardServiceStrategies = new ConcurrentHashMap<>();

	private RateLimiter apiGlobalLimiter;
	private RateLimiter rateLimiter;
	@Autowired(required = false)
	private StringRedisTemplate stringRedisTemplate;
	@Autowired(required = false)
	private DiscoveryClient discoveryClient;
	@Autowired
	private SystemApi systemApi;

	@Override
	public Object preHandle(ServerWebExchange exchange, BizSystemModule module, Builder requestBuilder) {
		ApiInfo api = RequestContextHelper.getCurrentApi(exchange);
		// 接口全局限流
		String hitKey = StringUtils.defaultIfBlank(api.getId(), api.getRoutedIdentifier());
		double permitsPerSecond = api.getPermitsPerSecond();
		if (permitsPerSecond > 0) {
			if (discoveryClient != null) {
				int podNums = CacheUtils.queryTryLevel1Cache(POD_NUM_CACHE_KEY, () -> {
					return discoveryClient.getInstances(GlobalContext.APPID).size();
				}, CacheExpires.IN_HALF_HOUR);
				//
				if (podNums > 1) {
					permitsPerSecond = new BigDecimal(permitsPerSecond)
							.divide(new BigDecimal(podNums), 3, RoundingMode.HALF_UP).doubleValue();
				}
			}
			if (!apiGlobalLimiter.tryAcquire(hitKey, permitsPerSecond)) {
				throw new MendmixBaseException(MainErrorType.TOO_MANY_REQUEST_LIMIT);
			}
		}
		//
		ApiInfo apiInfo = RequestContextHelper.getCurrentApi(exchange);
		if(apiInfo == null)return null;
		//优先动态规则
		RatelimitStrategy strategy = findActiveRuleByUri(module.getServiceId(), apiInfo.getIdentifier());
		if(strategy == null)return null;
		hitKey = strategy.resolveHitKey(exchange,module.getServiceId(),apiInfo);
		if(!rateLimiter.tryAcquire(hitKey, strategy.getPermitsPerSecond())) {
			throw new MendmixBaseException(MainErrorType.TOO_MANY_REQUEST_LIMIT);
		}
		
		return null;
	}

	@Override
	public String postHandle(ServerWebExchange exchange, boolean successed, String responseBody) {
		return null;
	}

	@Override
	public int ordered() {
		return 10;
	}
	
	@Override
	public void onStarted() {
		apiGlobalLimiter = new LocalRatelimiter();
		if(stringRedisTemplate == null) {
			logger.warn("<stringRedisTemplate> NOT FOUND!!");
			rateLimiter = new LocalRatelimiter();
		}else {
			rateLimiter = new RedisRatelimiter(stringRedisTemplate);
		}
		//加载规则
		initRatelimitRules();
	}
	
	public List<RatelimitStrategy> findRatelimitStrategies(String serviceId){
		List<RatelimitStrategy> rules = new ArrayList<>();
		if(serviceStrategies.containsKey(serviceId)) {
			rules.addAll(serviceStrategies.get(serviceId).values());
		}
		if(wildcardServiceStrategies.containsKey(serviceId)) {
			rules.addAll(wildcardServiceStrategies.get(serviceId).values());
		}
		return rules;
	}
	
    public String saveRatelimitStrategy(RatelimitStrategy strategy) {
    	String serviceId = strategy.getServiceId();
    	if(StringUtils.isAnyBlank(serviceId,strategy.getHitType(),strategy.getUriKey())) {
    		return null;
    	}
    	
    	BizSystemModule module = ProxySystemInfoHub.getModuleByServiceId(strategy.getServiceId());
    	if(module == null) {
    		throw new MendmixBaseException("服务模块["+serviceId+"]不存在");
    	}
    	
        if(StringUtils.isBlank(strategy.getId())) {
        	strategy.setId(DigestUtils.md5(serviceId + strategy.getHitType() + strategy.getUriKey()));
    	}
        //根据id找到原来的规则
        removeRatelimitStrategy(serviceId, strategy.getId());
    	
    	Map<String, RatelimitStrategy> uriRuleMapping = serviceStrategies.get(serviceId);
    	if(uriRuleMapping == null) {
    		serviceStrategies.put(serviceId, uriRuleMapping = new HashMap<>());
    	}
    	
    	Pattern pattern = strategy.getUriPattern();
		if(pattern == null) {
			uriRuleMapping.put(strategy.getUriKey(), strategy);
		}else {
			Map<Pattern, RatelimitStrategy> patternRuleMapping = wildcardServiceStrategies.get(serviceId);
		    if(patternRuleMapping == null) {
		    	wildcardServiceStrategies.put(serviceId, patternRuleMapping = new HashMap<>());
		    }
		    patternRuleMapping.put(pattern, strategy);
		}
    	
    	return strategy.getId();
	}
    
    public void removeRatelimitStrategy(String serviceId,String id) {
    	RatelimitStrategy strategy = findActiveRuleById(serviceId, id);
        if(strategy != null) {
        	if(strategy.getUriPattern() == null) {
        		serviceStrategies.get(strategy.getServiceId()).remove(strategy.getUriKey());
        	}else {
        		wildcardServiceStrategies.get(strategy.getServiceId()).remove(strategy.getUriPattern());
        	}
        }
    }
    
    public void removeRatelimitStrategy(RatelimitStrategy strategy) {
    	if(StringUtils.isBlank(strategy.getId())) {
        	strategy.setId(DigestUtils.md5(strategy.getServiceId() + strategy.getHitType() + strategy.getUriKey()));
    	}
    	removeRatelimitStrategy(strategy.getServiceId(), strategy.getId());
    	
    }
    
    private RatelimitStrategy findActiveRuleById(String serviceId,String id) {
    	RatelimitStrategy rule = null;
    	if(serviceStrategies.containsKey(serviceId)) {
    		rule = serviceStrategies.get(serviceId).values().stream().filter(o -> {
    			return o.getId().equals(id);
    		}).findFirst().orElse(null);
    	}
    	
    	if(rule == null && wildcardServiceStrategies.containsKey(serviceId)) {
    		rule = wildcardServiceStrategies.get(serviceId).values().stream().filter(o -> {
    			return o.getId().equals(id);
    		}).findFirst().orElse(null);
    	}
    	return rule;
    }
    
    private RatelimitStrategy findActiveRuleByUri(String serviceId,String uri) {
    	RatelimitStrategy rule = null;
    	if(serviceStrategies.containsKey(serviceId)) {
    		rule = serviceStrategies.get(serviceId).values().stream().filter(o -> {
    			return o.getUriKey().equals(uri);
    		}).findFirst().orElse(null);
    	}
    	
    	if(rule == null && wildcardServiceStrategies.containsKey(serviceId)) {
    		rule = wildcardServiceStrategies.get(serviceId).values().stream().filter(o -> {
    			return o.getUriPattern().matcher(uri).matches();
    		}).findFirst().orElse(null);
    	}
    	
    	if(rule == null && serviceStrategies.containsKey(serviceId)) {
    		rule = serviceStrategies.get(serviceId).get(ANY_URI_WILDCARD);
    	}
    	return rule;
    }

	private void initRatelimitRules() {
		//静态规则
		List<RatelimitStrategy> strategies = ResourceUtils.getConfigObjects("mendmix-cloud.ratelimit.rule", RatelimitStrategy.class);
	    int index = 1;
		for (RatelimitStrategy strategy : strategies) {
	    	strategy.setId(GatewayConstants.STATIC_RULE_ID_PREFIX + index);
	    	try {				
				saveRatelimitStrategy(strategy);
			} catch (Exception e) {
				logger.warn(">>initStaticRule:{},error:{}",strategy.getUriKey(),e.getMessage());
			}
	    	index++;
		}
		//
		List<BizSystem> systems = ProxySystemInfoHub.getSystems();
		for (BizSystem system : systems) {
			if(StringUtils.isBlank(system.getId()))continue;
			strategies = systemApi.getSystemRatelimitStrategies(system.getId());
			if(strategies == null)continue;
			for (RatelimitStrategy strategy : strategies) {
		    	try {				
					saveRatelimitStrategy(strategy);
				} catch (Exception e) {
					logger.warn(">>initRemoteRule:{},error:{}",strategy.getUriKey(),e.getMessage());
				}
			}
		}
		
	}

}
